from urlextract import URLExtract
import requests
from io import BytesIO
from cnocr import CnOcr
from PIL import Image
from DFA import DFAFilter
import sys


headers = {
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 '
                          '(KHTML, like Gecko) Chrome/113.0.0.0 Safari/537.36 Edg/113.0.1774.50'
        }


# ocr提取文字
def get_text(filepath: str):
    # img_fp = filepath # 这里得用相对路径
    req = requests.get(url=filepath, headers=headers)
    img_fp = Image.open(BytesIO(req.content))
    ocr = CnOcr()  # 所有参数都使用默认值
    out = ocr.ocr(img_fp=img_fp)

    text_list = []
    for i in out:
        if i['text'] != '':
            text_list.append(i['text'])
    return text_list


# 提取文档中的url
def get_url(markdown: str) -> list:
    extractor = URLExtract()
    urls = extractor.find_urls(markdown)
    return urls


# 根据url下载图片
def ocr_image(image_url: str):
    image_url = image_url.lower()
    # splits = image_url.split('/')
    # file_extension = splits[-1].split('.')[-1]  # 获取文件扩展名
    # 判断文件后缀
    endswith = ["webp", "bmp", "pcx", "tif", "gif", "jpeg", "tga", "exif", "fpx", "svg", "psd", "cdr", "pcd", "dxf",
                "ufo", "eps", "ai", "png", 'jpg', 'hdri', 'raw', 'wmf', 'flic', 'emf', 'ico']
    endswith_tuple = tuple(endswith)
    bool_endswith = image_url.endswith(endswith_tuple)
    if bool_endswith:
        # 下载图片
        # filename = int(round(datetime.now().timestamp(),-1)) # 用时间戳作为文件名
        # filepath = f'./images/{filename}.{file_extension}'
        # wget.download(image_url, filepath, bar=None)
        # req = requests.get(url=image_url, headers=headers)
        # open(filepath, 'wb').write(req.content)

        content = get_text(image_url)
        return check_image_sensitive(content)
        # 判断是否是图片
        # with open(image_url, mode='rb') as f:
        #     file_kind = filetype.guess(f)
        #     if file_kind is not None:
        #         # ocr提取文字
        #         # print('是图片')


# 检查是否有敏感词
def check_image_sensitive(content: list):
    text = ''.join(content)  # 将ocr取出的文字列表转为字符串
    gfw = DFAFilter()
    path = 'D:\\Study\\lsnu\\weiy-community-2023\\python\\files\\sensitive_words.txt'
    gfw.parse(path)
    result = gfw.filter(text)

    # print(text + ':' + result)
    return str({text:result})
    # time2 = time.time()
    # print('总共耗时：' + str(time2 - time1) + 's')


if __name__ == '__main__':

    # file = open('./files/test.md', encoding='utf-8')
    # doc = file.read()
    doc = sys.argv[1]
    url_list = get_url(doc)
    if len(url_list) != 0:
        result = ''
        for url in url_list:
            print(ocr_image(url))
        # print(result)
    else:
        print("'':''")

